from flask import Flask, request, jsonify, send_file,send_from_directory
import os,uuid,redis,shutil,json,zipfile

from flask.views import MethodView
from werkzeug.utils import secure_filename
# 引入异步任务
from tasks import async_video_processing, test_task, train_yolo, evaluate_model
# 数据集验证器
from selfutils.dataset_validator import DatasetValidator
# 视图模型
from appModels import Dataset, Model, db,cors

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'tmp_uploads' # 临时上传目录
app.config['RESULT_FOLDER'] = 'results' # 训练结果目录
app.config['MAX_CONTENT_LENGTH'] = 2 * 1024 * 1024 * 1024  # 文件大小限制2GB
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///datasets.db' # 数据库URI,/instance/datasets.db
app.config['DATASETS_UPLOAD_FOLDER'] = 'datasets' # 数据集上传目录

cors.init_app(app) # 初始化跨域
db.init_app(app) # 初始化数据库

"""
    数据库初始化：
flask shell
>>> from app import db
>>> db.create_all()
或者
>>> flask create_db
"""

# 初始化数据库
@app.cli.command()
def create_db():
    db.drop_all()
    with app.app_context():
        db.create_all()
        print("数据库初始化完成")

"""
    数据集管理
"""
# 获取数据集列表
@app.route('/datasets/info', methods=['GET'])
def get_datasets():
    try:
        page = request.args.get('page', 1, type=int)
        per_page = request.args.get('limit', 10, type=int)
        name_filter = request.args.get('name', '').strip()
        status_filter = request.args.get('status', '')

        # 参数验证
        if page < 1 or per_page < 1:
            raise ValueError('分页参数必须大于0')

        # 构建查询
        query = Dataset.query.order_by(Dataset.created_time.desc())
        # 应用名称筛选
        if name_filter:
            query = query.filter(Dataset.name.ilike(f'%{name_filter}%'))
        # 应用状态筛选
        if status_filter:
            query = query.filter(Dataset.status == status_filter)

        # 执行分页查询
        pagination = query.paginate(page=page, per_page=per_page, error_out=False)
        if not pagination.items:
            return jsonify({
                'code': 200,
                'message': 'success',
                'data': {
                    'items': [],
                    'total': 0
                }
            })

        # 构建响应数据
        datasets = [{
            'id': d.id,
            'name': d.name,
            'description': d.description,
            'train_samples': d.train_samples,
            'val_samples': d.val_samples,
            'num_classes': d.num_classes,
            'status': d.status,
            'created_time': d.created_time.isoformat()
        } for d in pagination.items]

        return jsonify({
            'code':200,
            'message':'success',
            'data':{
                'items': datasets,
                'total': pagination.total,
                'pages': pagination.pages,
                'current_page': page
            }
        })

    except ValueError as ve:
        return jsonify({'error': '参数错误', 'details': str(ve)}), 400
    except Exception as e:
        app.logger.error(f'获取数据集失败: {str(e)}')
        return jsonify({'error': '服务器内部错误'}), 500

# 数据集创建接口
@app.route('/datasets/create', methods=['POST'])
def create_dataset():
    # 验证请求数据
    if 'file' not in request.files:
        return jsonify({'error': 'No file uploaded'}), 400

    file = request.files['file']
    if not file.filename.endswith('.zip'):
        return jsonify({'error': 'Only ZIP files are allowed'}), 400

    # 获取表单数据
    dataset_name = request.form.get('name', 'untitled')
    description = request.form.get('description', '')

    # 创建数据集目录，此处 secure_filename(dataset_name)，不能只有中文，否则会报错
    dataset_dir = os.path.join(app.config['DATASETS_UPLOAD_FOLDER'], secure_filename(dataset_name))
    if os.path.exists(dataset_dir):
        return jsonify({'error': 'Dataset already exists'}), 409
    os.makedirs(dataset_dir)

    try:
        # 保存并解压ZIP文件
        zip_path = os.path.join(dataset_dir, file.filename)
        file.save(zip_path)

        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            # 自动处理嵌套目录
            root_dir = None
            for name in zip_ref.namelist():
                if '/' in name:
                    root_dir = name.split('/')[0]
                    break
            temp_dir = os.path.join(dataset_dir, 'temp')
            zip_ref.extractall(temp_dir)
            actual_dir = os.path.join(temp_dir, root_dir) if root_dir else temp_dir
            # 移动文件到目标目录
            for item in os.listdir(actual_dir):
                src = os.path.join(actual_dir, item)
                dst = os.path.join(dataset_dir, item)
                if os.path.exists(dst):
                    shutil.rmtree(dst)
                shutil.move(src, dst)
            shutil.rmtree(temp_dir)

        # 仅保留基础验证
        if not os.path.exists(os.path.join(dataset_dir, 'images')) or \
                not os.path.exists(os.path.join(dataset_dir, 'labels')):
            raise ValueError('数据集必须包含images和labels目录')

        # 保存到数据库（初始状态为未验证）
        new_dataset = Dataset(
            name=dataset_name,
            description=description,
            path=dataset_dir,
            status='unverified',
            train_samples=0,
            val_samples=0
        )
        db.session.add(new_dataset)
        db.session.commit()

        return jsonify({
            'code': 200,
            'data': {
                'id': new_dataset.id,
                'verify_url': f'/datasets/{new_dataset.id}/verify'
            }
        })

    except Exception as e:
        shutil.rmtree(dataset_dir, ignore_errors=True)
        return jsonify({
            'code': 400,
            'error': '创建失败',
            'details': str(e)
        }), 400

# 更新数据集信息
@app.route('/datasets/<int:dataset_id>', methods=['PUT'])
def update_dataset(dataset_id):
    dataset = Dataset.query.get_or_404(dataset_id)
    data = request.json

    try:
        if 'name' in data:
            dataset.name = data['name']
        if 'description' in data:
            dataset.description = data['description']

        db.session.commit()
        return jsonify({
            'code':200,
            'message': 'Dataset updated successfully'}
        )

    except Exception as e:
        db.session.rollback()
        return jsonify({'error': str(e)}), 400

# 删除数据集
@app.route('/datasets/<int:dataset_id>', methods=['DELETE'])
def delete_dataset(dataset_id):
    dataset = Dataset.query.get_or_404(dataset_id)

    try:
        # 删除文件
        if os.path.exists(dataset.path):
            shutil.rmtree(dataset.path)

        # 删除数据库记录
        db.session.delete(dataset)
        db.session.commit()
        return jsonify({
            'code':200,
            'message': 'Dataset deleted successfully'
        })

    except Exception as e:
        return jsonify({'error': str(e)}), 500

# 下载数据集
@app.route('/datasets/<int:dataset_id>/download')
def download_dataset(dataset_id):
    dataset = Dataset.query.get_or_404(dataset_id)
    zip_path = os.path.join(dataset.path, f'{dataset.name}.zip')

    try:
        # 创建ZIP文件
        with zipfile.ZipFile(zip_path, 'w') as zipf:
            for root, dirs, files in os.walk(dataset.path):
                for file in files:
                    if file.endswith('.zip'):
                        continue
                    file_path = os.path.join(root, file)
                    zipf.write(file_path, os.path.relpath(file_path, dataset.path))

        # 成功时直接返回文件流
        # return send_file(zip_path, as_attachment=True)
        return send_file(
            zip_path,
            as_attachment=True,
            download_name=f'{dataset.name}.zip',  # 关键修改！添加文件名参数
            mimetype='application/zip'
        )

    except Exception as e:
        # 错误时返回JSON响应
        return jsonify({
            'code': 500,
            'error': str(e)
        }), 500

# 数据集验证接口
@app.route('/datasets/<int:dataset_id>/verify', methods=['POST'])
def verify_dataset(dataset_id):
    dataset = Dataset.query.get_or_404(dataset_id)
    try:
        # 执行验证逻辑
        DatasetValidator.validate_structure(dataset.path) # 验证目录结构
        yaml_path, num_classes = DatasetValidator.validate_yaml(dataset.path) # 获取类别数
        counts = DatasetValidator.validate_pairs(dataset.path) # 获取训练和验证样本数

        # 更新数据集状态
        dataset.yaml_path = yaml_path
        dataset.num_classes = num_classes  # 存储类别数
        dataset.train_samples = counts.get('train_samples', 0)
        dataset.val_samples = counts.get('val_samples', 0)
        dataset.status = 'verified' if dataset.train_samples > 0 and dataset.val_samples > 0 else 'invalid'
        db.session.commit()

        return jsonify({
            'code': 200,
            'message': 'Dataset validated successfully',
            'data': {
                'status': dataset.status,
                'train_samples': dataset.train_samples,
                'val_samples': dataset.val_samples,
                'num_classes': dataset.num_classes,
            }
        })

    except ValueError as e:
        dataset.status = 'invalid'
        db.session.commit()
        return jsonify({
            'code': 400,
            'error': '验证失败',
            'details': str(e)
        }), 400

    except Exception as e:
        return jsonify({
            'code': 500,
            'error': '服务器错误',
            'details': str(e)
        }), 500


"""
    模型训练和管理Api
"""
# class ModelApi(MethodView):
#     def get(self, model_id):
#         if not model_id:
#             models = Model.query.all()
# model_api = ModelApi.as_view('model_api')
# app.add_url_rule('/models/info', view_func=model_api, methods=['GET', ])

@app.route('/models/info', methods=['GET'])
def get_models():
    page = request.args.get('page', 1, type=int)
    page_size = request.args.get('page_size', 10, type=int)
    model_name = request.args.get('model_name', '')
    model_type = request.args.get('model_type', '')

    # 分页参数校验
    if page < 1 or page_size < 1 or page_size > 100:
        return jsonify({'code': 400, 'message': 'Invalid pagination parameters'}), 400

    query = Model.query
    if model_name:
        query = query.filter(Model.name.ilike(f'%{model_name}%'))
    if model_type:
        query = query.filter(Model.model_type == model_type)  # 精确匹配

    pagination = query.paginate(page=page, per_page=page_size, error_out=False)
    return jsonify({
        'code': 200,
        'message': 'success',
        'data': {
            'items': [model.to_dict() for model in pagination.items],
            'total': pagination.total,
            'page': pagination.page,
            'pages': pagination.pages
        }
    })

@app.route('/models/create', methods=['POST'])
def create_model():
    data = request.get_json()
    if not data:
        return jsonify({'code': 400, 'message': 'No JSON data provided'}), 400

    required_fields = ['name', 'dataset_id', 'model_type']
    for field in required_fields:
        if field not in data:
            return jsonify({'code': 400, 'message': f'Missing field: {field}'}), 400

    try:
        model = Model(
            name=data['name'],
            description=data.get('description', ''),
            dataset_id=data['dataset_id'],
            model_type=data['model_type']
        )
        db.session.add(model)
        db.session.commit()
        return jsonify({
            'code': 200,
            'message': 'success',
            'data': model.to_dict()
        }), 200
    except Exception as e:
        db.session.rollback()
        return jsonify({'code': 500, 'message': str(e)}), 500

@app.route('/model/train', methods=['POST'])
def model_train_training():
    data = request.json
    task = train_yolo.apply_async(args=[data['dataset_id'], data['config']])
    return jsonify({"task_id": task.id}), 202

@app.route('/evaluate', methods=['POST'])
def model_evaluation():
    data = request.json
    task = evaluate_model.apply_async(args=[data['model_path'], data['dataset_id']])
    return jsonify({"task_id": task.id}), 202


"""
    视频处理
    上传视频
    查看任务状态
    下载处理后的视频
    查看原始视频
    查看处理后的视频
    行为检测结果可视化接口
"""
@app.route('/upload', methods=['POST'])
def upload_video():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400

    # 生成Celery任务ID
    async_task_id = str(uuid.uuid4())  # 提前生成用于目录创建

    # 使用Celery任务ID创建任务目录
    task_folder = os.path.join(app.config['UPLOAD_FOLDER'], async_task_id)
    os.makedirs(task_folder, exist_ok=True)

    # 保存上传文件
    filename = secure_filename(file.filename)
    input_path = os.path.join(task_folder, filename)
    output_path = os.path.join(app.config['RESULT_FOLDER'], f'processed_{async_task_id}.mp4')
    file.save(input_path)

    # # 获取用户选择的模型参数
    # yolo_model = request.form.get('yolo_model', 'yolov5l6')
    # slowfast_model = request.form.get('slowfast_model', 'slowfast_r50_detection')
    #
    # config_args = {
    #     'yolo_model': yolo_model,
    #     'slowfast_model': slowfast_model,
    #     # ...其他原有参数...
    # }

    # 启动异步任务时传递完整参数
    config_args = {
        'imsize': 640,
        'conf': 0.4,
        'iou': 0.4,
        'device': 'cuda',  # 或 'cpu'
        'show': False,
        'classes': None  # 必须显式传递所有参数
    }

    # 调用脚本处理视频
    async_result = async_video_processing.apply_async(
        args=(async_task_id, input_path, output_path, config_args),
        task_id=async_task_id  # 显式指定任务ID
    )
    async_task_id = async_result.id  # 获取任务ID

    return jsonify({
        'task_id': async_task_id,
        'status_url': f'/status/{async_task_id}',
        'download_url': f'/download/{async_task_id}',
        'original_url': f'/files/{async_task_id}/{filename}',  # 新增原始文件地址
        'processed_url': f'/processed_files/processed_{async_task_id}.mp4',  # 新增处理结果地址
    }), 202

@app.route('/status/<task_id>')
def get_status(task_id):
    # 获取任务状态
    task = async_video_processing.AsyncResult(task_id)

    response = {
        'task_id': task_id,
        'celery_status': task.state,  # 保留原始状态用于调试
        'status': 'pending',
        'progress': 0
    }

    if task.state == 'PROGRESS':
        response.update({
            'status': 'processing',
            'progress': task.info.get('progress', 0),
            'current': task.info.get('current', 0),
            'total': task.info.get('total', 0)
        })
    elif task.state == 'SUCCESS':
        response.update({
            'status': task.result['status'],  # 从result获取业务状态
            'progress': 100,
            'result_path': task.result.get('result_path', '')
        })
    elif task.state == 'FAILURE':
        response.update({
            'status': 'failed',
            'error': task.result.get('error', 'Unknown error')
        })

    return jsonify(response)

@app.route('/download/<task_id>')
def download_video(task_id):
    task_folder = os.path.join(app.config['UPLOAD_FOLDER'], task_id)
    # result_file = os.path.join(app.config['RESULT_FOLDER'], f'processed_{os.listdir(task_folder)[0]}')
    # 处理结果路径
    result_file = os.path.join(
        app.config['RESULT_FOLDER'],
        f'processed_{task_id}.mp4'
    )
    if not os.path.exists(result_file):
        return jsonify({'error': 'File not ready'}), 404

    return send_file(result_file, as_attachment=True)

@app.route('/files/<task_id>/<filename>')
def serve_file(task_id, filename):
    """通用文件访问接口"""
    return send_from_directory(
        os.path.join(app.config['UPLOAD_FOLDER'], task_id),
        filename,
        mimetype='video/mp4'
    )

@app.route('/processed_files/<filename>')
def serve_processed_file(filename):
    """处理结果文件访问接口"""
    return send_from_directory(
        app.config['RESULT_FOLDER'],
        filename,
        mimetype='video/mp4'
    )


@app.route('/original_video/<task_id>')
def get_original_video(task_id):
    task_folder = os.path.join(app.config['UPLOAD_FOLDER'], task_id)
    original_file = os.path.join(task_folder, os.listdir(task_folder)[0])

    if not os.path.exists(original_file):
        return jsonify({'error': 'Original file not found'}), 404

    return send_file(original_file, as_attachment=False)  # 不作为附件下载

@app.route('/processed_video/<task_id>')
def get_processed_video(task_id):
    result_file = os.path.join(
        app.config['RESULT_FOLDER'],
        f'processed_{task_id}.mp4'
    )

    if not os.path.exists(result_file):
        return jsonify({'error': 'File not ready'}), 404

    # 强制设置视频MIME类型和响应头
    response = send_file(
        result_file,
        mimetype='video/mp4',
        as_attachment=False,# 不作为附件下载
        conditional=True
    )

    # 添加流媒体播放支持
    response.headers['Accept-Ranges'] = 'bytes'
    response.headers['Cache-Control'] = 'no-cache'
    return response

# 新增行为检测结果存储和分析接口
@app.route('/analysis/<task_id>')
def get_analysis_data(task_id):
    """获取行为分析结果数据"""
    try:
        # 假设每个任务会生成一个对应的JSON分析文件
        analysis_path = os.path.join(app.config['RESULT_FOLDER'], f'analysis_{task_id}.json')

        if not os.path.exists(analysis_path):
            return jsonify({'error': 'Analysis data not ready'}), 404

        with open(analysis_path, 'r') as f:
            data = json.load(f)

        return jsonify({
            'task_id': task_id,
            'status': 'completed',
            'data': data
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500


# 添加Redis健康检查接口
@app.route('/redis-check')
def redis_check():
    try:
        r = redis.Redis()
        return jsonify({
            "status": "success",
            "ping": r.ping(),
            "info": {
                "version": r.info().get("redis_version"),
                "memory_used": r.info().get("used_memory_human")
            }
        })
    except Exception as e:
        return jsonify({"status": "error", "error": str(e)})

@app.route('/test')
def test():
    test_task.apply_async()
    return jsonify({"message": "测试任务已触发"})


if __name__ == '__main__':
    # 初始化目录
    os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
    os.makedirs(app.config['RESULT_FOLDER'], exist_ok=True)
    # with app.app_context():
    #     db.create_all()
    app.run(host='0.0.0.0', port=5000,debug=True)


